# -*- coding: utf-8 -*-
# @Time    : 2022/7/24 20:16
# @Author  : Aweo0419
import onnx
import onnxruntime as ort
import numpy as np
from torch.autograd import Variable
from torchvision import transforms
import cv2
import torch

def img_read(img):
    # img = cv2.imread(img, cv2.IMREAD_COLOR)
    img = cv2.resize(img, (32, 32))
    transform = transforms.Compose(
        [transforms.ToTensor(),
         transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
    img = Variable(torch.unsqueeze(transform(img), dim=0).float(), requires_grad=False)
    return img


def to_numpy(tensor):
    return tensor.detach().cpu().numpy().astype(np.float32)


def onnx_runtime():
    t_img = cv2.imread("img_data/nag_img/1.png")
    test_ar = img_read(t_img)

    sess = ort.InferenceSession('test.onnx')
    # onnx_model = onnx.load("test.onnx")
    # onnx.checker.check_model(onnx_model)
    # print(1)
    input_name = sess.get_inputs()[0].name

    output_name = sess.get_outputs()[0].name

    pred_onnx = sess.run([output_name], {input_name:to_numpy(test_ar)})

    print("outputs:")
    print(np.array(pred_onnx))


onnx_runtime()
